#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
该模块通过调用 LLM 模型，根据给定的研究主题和文本内容查找相关的语句引用，
并支持并发处理文本段。所有功能均封装在 CitationProcessor 类中。
"""

import asyncio
import logging
from asyncio import Semaphore
from pathlib import Path
from typing import List, Optional

from jinja2 import Environment
from pyaml_env import parse_config
import json_repair

from research_agent.core.config import Config
from research_agent.core.general_llm import LLM
from research_agent.core.query import Query
from research_agent.core.utils import tokenize_sentences

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


class FindStatementCitation:
    """
    通过调用 LLM 模型，根据给定的研究主题和文本内容查找相关的语句引用，
    并提供对文本段并发处理的功能。
    """

    def __init__(self, base_path: Optional[str] = None):
        """
        初始化 FindStatementCitation 类，加载 LLM 模型配置和提示模板，并设置并发控制信号量。

        :param base_path: 提示模板所在的基础路径，默认为当前文件所在目录下的 "prompts" 文件夹
        """
        # 加载配置文件并获取默认模型配置
        self.configs = parse_config(Config.YAML_CONFIG)
        self.llm = LLM(config=self.configs[Config.DEFAULT_MODEL])
        self.query = Query()

        # 确定模板文件路径
        if base_path is None:
            base_path = Path(__file__).parent / "prompts"
        else:
            base_path = Path(base_path)

        prompt_file = base_path / "find_statements.jinja"
        try:
            with open(prompt_file, "r", encoding="utf-8") as f:
                template_content = f.read()
        except Exception as e:
            logger.error(f"加载提示模板文件失败：{prompt_file}，错误信息：{e}")
            raise

        # 使用 Jinja2 加载模板
        self.prompt_template = Environment().from_string(template_content)

    def _prepare_prompt_messages(self, topic: str, section: str) -> List[dict]:
        """
        准备生成查找语句引用所需的提示消息。

        :param topic: 研究主题
        :param section: 文本内容
        :return: 包含系统和用户提示信息的字典列表
        """
        system_prompt = self.prompt_template.render(role="system", topic=topic)
        user_prompt = self.prompt_template.render(
            role="user",
            survey_draft=tokenize_sentences(section),
            topic=topic
        )
        return [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_prompt},
        ]

    async def find_statement_citation(self, topic: str, section: str, max_retries: int = 3) -> List[str]:
        """
        根据研究主题和文本内容调用 LLM 生成回答，提取引用语句。

        :param topic: 研究主题
        :param section: 预处理后的文本内容
        :param max_retries: 最大重试次数
        :return: 模型返回的引用语句列表
        """
        prompt_messages = self._prepare_prompt_messages(topic, section)
        for attempt in range(max_retries):
            try:
                response = await self.llm.completion(prompt_messages)
                response_data = json_repair.loads(response)
                if isinstance(response_data, dict) and "statements" in response_data:
                    return response_data["statements"]
                if isinstance(response_data, list) and len(response_data) > 0 and isinstance(response_data[0], dict) and "statement_abstract" in response_data[0] and "keywords" in response_data[0] and "evidence_spans" in response_data[0]:
                    return response_data
            except Exception as e:
                logger.error(f"调用 LLM 模型时出错：{e}，尝试次数：{attempt + 1}")
                if attempt < max_retries - 1:
                    await asyncio.sleep(0.5)
                else:
                    logger.error("达到最大重试次数，操作失败。")
                    return []

    async def process_section(self, section: str, topic: str) -> Optional[List[str]]:
        """
        异步处理单个文本段：
          1. 对文本段进行句子分割和格式化；
          2. 调用 LLM 模型查找引用语句。

        :param section: 原始文本段
        :param topic: 研究主题
        :return: 模型返回的引用语句列表，出现异常时返回 None
        """
        try:
            citations = await self.find_statement_citation(topic=topic, section=section)
            return citations
        except Exception as e:
            logger.error(f"处理文本段时出错：{e}")
            return None

    async def process_all_sections(self, sections: List[str], topic: str) -> List[List[str]]:
        """
        并发处理所有文本段，保持原有顺序并返回每个文本段对应的引用语句结果。

        :param sections: 文本段列表
        :param topic: 研究主题
        :return: 每个文本段返回的引用语句列表集合（过滤掉处理失败的结果）
        """
        tasks = [self.process_section(section, topic) for section in sections]
        results = await asyncio.gather(*tasks, return_exceptions=True)
        return [result for result in results if result is not None]


# 示例入口
if __name__ == "__main__":
    async def main():
        topic = "人工智能在医疗领域的应用"
        sections = [
            "人工智能可以帮助医生更快地诊断疾病，提供个性化治疗方案。",
            "利用深度学习算法，可以从大量医学影像中提取关键特征。",
        ]
        processor = FindStatementCitation()
        citations_list = await processor.process_all_sections(sections, topic)
        for idx, citations in enumerate(citations_list, start=1):
            print(f"文本段 {idx} 的引用语句：")
            for citation in citations:
                print(citation)
            print("-" * 40)

    asyncio.run(main())
